import pandas as pd
import numpy as np
import tushare as ts
from pylab import mpl
from plot import MplVisualIf
app = MplVisualIf()
mpl.rcParams['font.sans-serif']=['SimHei']
mpl.rcParams['axes.unicode_minus']=False
# 可视化类（画各种图，各种花里胡哨）
class visualize:
    def ts_pro(self):
        import tushare as ts
        token = 'ea7e78a4c2d3b80b1566d22b24c6d09f9dfab662ccfe8a2cd1050d67'
        ts.set_token(token)
        pro = ts.pro_api(token)
        return pro

    def get_data(self, code, start='20000101', end=''):
        try:
            data = ts.pro_bar(ts_code=code, adj='qfq', start_date=start, end_date=end)
            data.index = pd.to_datetime(data.trade_date)
            data = data.sort_index()
        except:
            print('获取不到数据，请检查代码是否有误')
        return data

    def get_code(self):
        dd = self.ts_pro().stock_basic()
        code = dd.ts_code.values
        name = dd.name.values
        stocks = dict(zip(name, code))
        return stocks

    # stocks = get_code()
    # K线图
    def draw_kline(self,data, title):
        layout_dict = {'figsize': (12, 6),
                       'index': data.index,
                       'draw_kind': {'ochl':
                                         {'open': data.open,
                                          'close': data.close,
                                          'high': data.high,
                                          'low': data.low
                                          }
                                     },
                       'title': title + '-日K线',
                       'ylabel': "收盘价",
                       }
        app.fig_output(**layout_dict)

    # 成交量
    def draw_volume(self,data, title):
        bar_red = np.where(data.open < data.close, data.vol, 0)
        bar_green = np.where(data.open > data.close, data.vol, 0)
        layout_dict = {'figsize': (14, 5),
                       'index': data.index,
                       'draw_kind': {'bar':
                                         {'bar_red': bar_red,
                                          'bar_green': bar_green,
                                          }
                                     },
                       'title': title + '-成交量',
                       'ylabel': "成交量"}
        app.fig_output(**layout_dict)

    # 移动均线
    def draw_sma(self,data, title):
        data['sma20'] = data.close.rolling(20).mean()
        data['sma30'] = data.close.rolling(30).mean()
        data['sma60'] = data.close.rolling(60).mean()

        layout_dict = {'figsize': (14, 5),
                       'index': data.index,
                       'draw_kind': {'line':
                                         {'sma20': data.sma20,
                                          'sma30': data.sma30,
                                          'sma60': data.sma60
                                          }
                                     },
                       'title': title + '-均线',
                       'ylabel': "收盘价",
                       'legend': 'best'}
        app.fig_output(**layout_dict)

    # 震荡类指标可视化
    def draw_kdj(self,data, title):
        low_list = data.low.rolling(9, min_periods=1).min()
        high_list = data.high.rolling(9, min_periods=1).max()
        rsv = (data.close - low_list) / (high_list - low_list) * 100
        data['K'] = rsv.ewm(com=2, adjust=False).mean()
        data['D'] = data['K'].ewm(com=2, adjust=False).mean()
        data['J'] = 3 * data['K'] - 2 * data['D']

        layout_dict = {'figsize': (14, 5),
                       'index': data.index,
                       'draw_kind': {'line':
                                         {'K': data.K,
                                          'D': data.D,
                                          'J': data.J
                                          }
                                     },
                       'title': title + '-KDJ',
                       'ylabel': "KDJ",
                       'legend': 'best'}
        app.fig_output(**layout_dict)

    # 趋势类指标MACD可视化
    def draw_macd(self,data, title):
        macd_dif = data.close.ewm(span=12, adjust=False).mean() - data.close.ewm(span=26, adjust=False).mean()
        macd_dea = macd_dif.ewm(span=9, adjust=False).mean()
        macd_bar = 2 * (macd_dif - macd_dea)
        bar_red = np.where(macd_bar > 0, macd_bar, 0)
        bar_green = np.where(macd_bar < 0, macd_bar, 0)
        layout_dict = {'figsize': (14, 5),
                       'index': data.index,
                       'draw_kind': {'bar':
                                         {'bar_red': bar_red,
                                          'bar_green': bar_green,
                                          },
                                     'line':
                                         {'macd_dif': macd_dif,
                                          'macd_dea': macd_dea}
                                     },
                       'title': title + '-MACD',
                       'ylabel': "MACD"}
        app.fig_output(**layout_dict)

    def draw_stocker(self):
        layout_dict = {'figsize': (12, 8),
                       'nrows': 4,
                       'ncols': 1,
                       'left': 0.07,
                       'bottom': 0.15,
                       'right': 0.99,
                       'top': 0.96,
                       'wspace': None,
                       'hspace': 0,
                       'height_ratios': [3.5, 1, 1, 1],
                       'subplots': ['kgraph', 'volgraph', 'kdjgraph', 'macdgraph']}
        subplots_dict = {'graph_fst': {'graph_name': 'kgraph',
                                       'graph_type': {'ochl': None,
                                                      'ma': [20, 30, 60, ]},
                                       'title': '日K线',
                                       'ylabel': '价格',
                                       'xticks': 15,
                                       'legend': 'best'
                                       },
                         'graph_sec': {'graph_name': 'volgraph',
                                       'graph_type': {'vol': None,
                                                      },
                                       'ylabel': '成交量',
                                       'xticks': 15,
                                       },
                         'graph_thd': {'graph_name': 'kdjgraph',
                                       'graph_type': {'kdj': None,
                                                      },
                                       'ylabel': 'KDJ',
                                       'xticks': 15,
                                       'legend': 'best'
                                       },
                         'graph_fth': {'graph_name': 'macdgraph',
                                       'graph_type': {'macd': None,
                                                      },
                                       'ylabel': 'MACD',
                                       'xlabel': '日期',
                                       'xticks': 15,
                                       'legend': 'best',
                                       'xticklabels': '%Y%m%d'
                                       },
                         }
        # draw_stock = MultiGraphIf(**layout_dict)
        # draw_stock.graph_run(stock_dat.copy(deep=True), **subplots_dict)
        #
        # draw_stock.graph_run(stock_dat.copy(deep=True), **subplots_dict)

    def show_detail(self,data,stock):
        self.draw_kline(data, stock)
        self.draw_volume(data, stock)  # 成交量
        self.draw_sma(data, stock)  # 移动均线
        self.draw_kdj(data, stock)  # 震荡类指标可视化
        self.draw_macd(data, stock)  # 趋势类指标MACD可视化
        self.draw_stocker()



